#include <string>
#include <set>
#include <vector>
#include <iostream>
#include <fstream>
#include <io.h>
#include <map>
#include <algorithm>
#include <cmath>
#include <iomanip>
using namespace std;

/*
Project:   Infomation Retrieval VSM
Author:    WCX
Time:      2022/11/9
*/

// number of docs
int docNum = 0;
// weights of different scopes
const double titleWeight = 0.5;
const double authorWeight = 0.2;
const double contentWeight = 0.3;
// Use this struct to record docs' score
struct DocScore {
    string fileName;
    double score;
}docScore[1000];
bool operator<(const DocScore& a, const DocScore& b) {
    if (a.score == b.score) {
        return a.fileName < b.fileName;
    }
    return a.score > b.score;
}
// Get paths and names of all test files
void getFiles(const string& dirPath, vector<string>& filePaths, vector<string>& fileNames) {
    string path;
    long hFile = 0; //file handle
    struct _finddata_t fileInfo; //file infomation
    if ((hFile = _findfirst(path.assign(dirPath).append("/*").c_str(), &fileInfo)) != 0) {
        do {
            string path_s = path.assign(dirPath).append("/").append(fileInfo.name);
            if (!strcmp(fileInfo.name, ".") || !strcmp(fileInfo.name, "..")) {
                continue;
            }
            string fileName = fileInfo.name;
            transform(fileName.begin(), fileName.end(), fileName.begin(), ::tolower); // Transform to lower case
            fileNames.push_back(fileName.substr(0, fileName.rfind("."))); // Remove ".txt"
            filePaths.push_back(path_s);
        } while (_findnext(hFile, &fileInfo) != -1);
        _findclose(hFile);
    }
}
// Tokenize string input
void tokenize(const string& str, vector<string>& tokens, const string& delimiters) {
    string::size_type lastPos = str.find_first_not_of(delimiters, 0);
    string::size_type pos = str.find_first_of(delimiters, lastPos);
    while (string::npos != pos || string::npos != lastPos) {
        string token = str.substr(lastPos, pos - lastPos);
        transform(token.begin(), token.end(), token.begin(), ::tolower); // Transform to lower case
        tokens.push_back(token);
        lastPos = str.find_first_not_of(delimiters, pos);
        pos = str.find_first_of(delimiters, lastPos);
    }
}
// Sort unique tokens
void sortUniqueTokens(vector<string>& tokens) {
    set<string> st(tokens.begin(), tokens.end());
    tokens.assign(st.begin(), st.end());
}
// Get df-Map
void getDfMap(vector<string>& tokens, map<string, int>& dfMap) {
    for (unsigned int i = 0; i < tokens.size(); i++) {
        if (dfMap.find(tokens[i]) == dfMap.end()) {
            dfMap[tokens[i]] = 1;
        }
        else {
            dfMap[tokens[i]]++;
        }
    }
}
// Get tf-Map
void getTfMap(vector<string>& tokens, vector<string>& currentTokens, map<string, int>& tfMap) {
    for (unsigned int i = 0; i < tokens.size(); i++) {
        tfMap[tokens[i]] = 0;
    }
    for (unsigned int i = 0; i < currentTokens.size(); i++) {
        if (tfMap.find(currentTokens[i]) != tfMap.end()) {
            tfMap[currentTokens[i]]++;
        }
    }
}
void getRetrievalTfMap(map<string, int>& dfMap, vector<string>& currentTokens, map<string, int>& tfMap) {
    for (unsigned int i = 0; i < currentTokens.size(); i++) {
        if (dfMap.find(currentTokens[i]) != dfMap.end()) {
            if (tfMap.find(currentTokens[i]) == tfMap.end()) {
                tfMap[currentTokens[i]] = 1;
            }
            else {
                tfMap[currentTokens[i]]++;
            }
        }
    }
}
// Calculate doc vector length
void getDocVecLenMap(map<string, double>& docVecLenMap, map<string, map<string, int>>& docTfMapMap, map<string, int>& dfMap) {
    double vecLen = 0;
    map<string, map<string, int>>::iterator docIt;
    for (docIt = docTfMapMap.begin(); docIt != docTfMapMap.end(); docIt++) {
        map<string, int>::iterator tokenIt;
        for (tokenIt = (docIt->second).begin(); tokenIt != (docIt->second).end(); tokenIt++) {
            int tf = tokenIt->second;
            int df = dfMap[tokenIt->first];
            double tfIdf = log(tf + 1) * (log(docNum) - log(df));
            vecLen += pow(tfIdf, 2);
        }
        vecLen = pow(vecLen, 0.5);
        docVecLenMap[docIt->first] = vecLen;
        vecLen = 0;
    }
}
// Calculate cosine
double getCosine(map<string, int>& tfMap, map<string, int>& dfMap, map<string, int>& retrievalTfMap, double docVecLen) {
    double cosine = 0;
    map<string, int>::iterator it;
    for (it = retrievalTfMap.begin(); it != retrievalTfMap.end(); it++) {
        int retrievalTf = it->second;
        int tf = tfMap[it->first];
        int df = dfMap[it->first];
        double tfIdfProduct = log(tf + 1) * (log(docNum) - log(df)) * log(retrievalTf + 1) * (log(docNum) - log(df));
        cosine += tfIdfProduct;
    }
    cosine /= docVecLen;
    return cosine;
}
// Calculate score
double getScore(double titleScore, double authorScore, double contentScore) {
    double score = titleScore * titleWeight + authorScore * authorWeight + contentScore * contentWeight;
    return score;
}
// [DEBUG] Print contents
void displayTokens(vector<string>& tokens) {
    vector<string>::iterator it;
    cout << "Vector Size : " << tokens.size() << endl;
    for (it = tokens.begin(); it != tokens.end(); it++) {
        cout << *it << endl;
    }
}
void displayDfMap(map<string, int>& mapI) {
    map<string, int>::iterator it;
    cout << "Map Size : " << mapI.size() << endl;
    for (it = mapI.begin(); it != mapI.end(); it++) {
        cout << it->first << " " << it->second << endl;
    }
}
void displayTfMap(map<string, map<string, int>>& mapI) {
    map<string, map<string, int>>::iterator it;
    cout << "Map Size : " << mapI.size() << endl;
    for (it = mapI.begin(); it != mapI.end(); it++) {
        cout << it->first << " : " << endl << "{" << endl;
        map<string, int>::iterator jt;
        for (jt = (it->second).begin(); jt != (it->second).end(); jt++) {
            cout << "\t" << jt->first << " " << jt->second << endl;
        }
        cout << "}" << endl;
    }
}
void displayRetrievalTfMap(map<string, int>& tfMap) {
    cout << "Map Size : " << tfMap.size() << endl;
    map<string, int>::iterator it;
    for (it = tfMap.begin(); it != tfMap.end(); it++) {
        cout << it->first << " " << it->second << endl;
    }
}
void displayDocRank() {
    cout << "\n======================= Retrieval Result =======================\n";
    for (int i = 0; i < docNum; i++) {
        cout << setw(50) << left << docScore[i].fileName << left << "(" << docScore[i].score << ")" << endl;
    }
}

int main(int argc, char* argv[])
{
    // test files
    vector<string>filePaths;
    vector<string>fileNames;
    // titles && authors && tokens
    vector<string>titles;
    vector<string>authors;
    vector<string>tokens;
    // token df-Maps
    map<string, int>titleDfMap;
    map<string, int>authorDfMap;
    map<string, int>tokenDfMap;
    // docs' tf-Maps
    map<string, map<string, int>> docTitleTfMapMap;
    map<string, map<string, int>> docAuthorTfMapMap;
    map<string, map<string, int>> docTokenTfMapMap;
    // doc vector-Maps
    map<string, double> docTitleVecLenMap;
    map<string, double> docAuthorVecLenMap;
    map<string, double> docTokenVecLenMap;
    // attributes
    const string testDir = "dataset";
    const string delimiters = " ,;.!?'\":";

    getFiles(testDir, filePaths, fileNames);
    docNum = fileNames.size();
    /* Get DfMaps */
    // Get titleDfMap
    for (unsigned int i = 0; i < fileNames.size(); i++) {
        vector<string>currentTitle;
        tokenize(fileNames[i], currentTitle, delimiters);
        tokenize(fileNames[i], titles, delimiters);
        sortUniqueTokens(currentTitle);
        getDfMap(currentTitle, titleDfMap);
    }
    sortUniqueTokens(titles);
    // Get authorDfMap && tokenDfMap
    ifstream fin;
    char buf[1024] = { 0 };
    for (unsigned int i = 0; i < filePaths.size(); i++) {
        vector<string>currentAuthor;
        vector<string>currentTokens;
        fin.open(filePaths[i], ios::in);
        if (!fin.is_open()) {
            std::cerr << "Cannot open file : " << filePaths[i] << endl;
        }
        // Get authors
        fin.getline(buf, sizeof(buf));
        string authorName = buf;
        tokenize(authorName.substr(7), currentAuthor, delimiters);
        tokenize(authorName.substr(7), authors, delimiters);
        sortUniqueTokens(currentAuthor);
        getDfMap(currentAuthor, authorDfMap);
        // Get tokens
        while (fin.getline(buf, sizeof(buf))) {
            tokenize(buf, currentTokens, delimiters);
            tokenize(buf, tokens, delimiters);
        }
        sortUniqueTokens(currentTokens);
        getDfMap(currentTokens, tokenDfMap);
        fin.close();
    }
    sortUniqueTokens(authors);
    sortUniqueTokens(tokens);
    //displayTokens(tokens);
    //displayDfMap(tokenDfMap);
    /* Get tfMaps */
    // Get titleTfMap
    for (unsigned int i = 0; i < fileNames.size(); i++) {
        vector<string>currentTitle;
        map<string, int>titleTfMap;
        tokenize(fileNames[i], currentTitle, delimiters);
        getTfMap(titles, currentTitle, titleTfMap);
        docTitleTfMapMap[fileNames[i]] = titleTfMap;
    }
    // Get authorTfMap && tokenTfMap
    for (unsigned int i = 0; i < filePaths.size(); i++) {
        vector<string>currentAuthor;
        vector<string>currentTokens;
        map<string, int>authorTfMap;
        map<string, int>tokenTfMap;
        fin.open(filePaths[i], ios::in);
        if (!fin.is_open()) {
            std::cerr << "Cannot open file : " << filePaths[i] << endl;
        }
        // Get authors
        fin.getline(buf, sizeof(buf));
        string authorName = buf;
        tokenize(authorName.substr(7), currentAuthor, delimiters);
        getTfMap(authors, currentAuthor, authorTfMap);
        docAuthorTfMapMap[fileNames[i]] = authorTfMap;
        // Get tokens
        while (fin.getline(buf, sizeof(buf))) {
            tokenize(buf, currentTokens, delimiters);
        }
        getTfMap(tokens, currentTokens, tokenTfMap);
        docTokenTfMapMap[fileNames[i]] = tokenTfMap;
        fin.close();
    }
    //displayTfMap(docTokenTfMapMap);
    getDocVecLenMap(docTitleVecLenMap, docTitleTfMapMap, titleDfMap);
    getDocVecLenMap(docAuthorVecLenMap, docAuthorTfMapMap, authorDfMap);
    getDocVecLenMap(docTokenVecLenMap, docTokenTfMapMap, tokenDfMap);
    
    /* Retrieval Proccess */
    // title && author && tokens
    vector<string>retrievalTitle;
    vector<string>retrievalAuthor;
    vector<string>retrievalTokens;
    // retrieval tf-Maps
    map<string, int>retrievalTitleTfMap;
    map<string, int>retrievalAuthorTfMap;
    map<string, int>retrievalTokenTfMap;
    bool retrieveTitle = false, retrieveAuthor = false, retrieveContent = false;
    char input[1024] = { 0 };
    cout << "Title >> ";
    gets_s(input);
    tokenize(input, retrievalTitle, delimiters);
    if (!retrievalTitle.empty()) {
        retrieveTitle = true;
        getRetrievalTfMap(titleDfMap, retrievalTitle, retrievalTitleTfMap);
        //displayRetrievalTfMap(retrievalTitleTfMap);
    }
    else {
        cout << "Ok, no retrival condition for the title scope.\n";
    }
    cout << "Author >> ";
    gets_s(input);
    tokenize(input, retrievalAuthor, delimiters);
    if (!retrievalAuthor.empty()) {
        retrieveAuthor = true;
        getRetrievalTfMap(authorDfMap, retrievalAuthor, retrievalAuthorTfMap);
        //displayRetrievalTfMap(retrievalAuthorTfMap);
    }
    else {
        cout << "Ok, no retrival condition for the author scope.\n";
    }
    cout << "Content >> ";
    gets_s(input);
    tokenize(input, retrievalTokens, delimiters);
    if (!retrievalTokens.empty()) {
        retrieveContent = true;
        getRetrievalTfMap(tokenDfMap, retrievalTokens, retrievalTokenTfMap);
        //displayRetrievalTfMap(retrievalTokenTfMap);
    }
    else {
        cout << "Ok, no retrival condition for the content scope.\n";
    }
    // Calculate scores and sort docs
    for (unsigned int i = 0; i < fileNames.size(); i++) {
        docScore[i].fileName = fileNames[i];
        double titleScore = 0, authorScore = 0, contentScore = 0;
        if (retrieveTitle) {
            titleScore = getCosine(docTitleTfMapMap[fileNames[i]], titleDfMap, retrievalTitleTfMap, docTitleVecLenMap[fileNames[i]]);
        }
        if (retrieveAuthor) {
            authorScore = getCosine(docAuthorTfMapMap[fileNames[i]], authorDfMap, retrievalAuthorTfMap, docAuthorVecLenMap[fileNames[i]]);
        }
        if (retrieveContent) {
            contentScore = getCosine(docTokenTfMapMap[fileNames[i]], tokenDfMap, retrievalTokenTfMap, docTokenVecLenMap[fileNames[i]]);
        }
        docScore[i].score = getScore(titleScore, authorScore, contentScore);
    }
    sort(docScore, docScore + docNum);
    displayDocRank();
    return 0;
}